#include <stdlib.h>
#include "cublas_v2.h"
#include <stdio.h>
#include "mkl_lapack.h"
#include "cuda.h"
#include "cuda_runtime_api.h"
#include "dept150_common.h"

#define BLOCK_SIZE  16

int LPPSADB_cu_multi_gpu_ck(int n, double* a, double* b, double *E, double *El, double *cond, double ea, double eb);
int LPPSADB_cu_multi_gpu_ck(int block_size, int n, double* a, double* b, double *E, double *El, double *cond, double ea, double eb);

int LPPSADB_cu_multi_gpu_ck(int n, double* a, double* b, double *E, double *El, double *cond, double ea, double eb)
{
	
	return LPPSADB_cu_multi_gpu_ck(BLOCK_SIZE, n, a, b, E, El, cond, ea, eb);
}

int LPPSADB_cu_multi_gpu_ck(int block_size, int n, double* a, double* b, double *E, double *El, double *cond, double ea, double eb)
{
#ifdef DETAIL_DEBUG
	printf("given matrix\n");
	print_matrix(n,n,a,n);
	printf("===");
#endif
	int gpu_n;
	checkCudaErrors(cudaGetDeviceCount(&gpu_n));
#ifdef DEBUG
	printf("Found %i devices\n", gpu_n);
#endif
	if(gpu_n < 1)
	{
		printf("Can not find any GPU\n");
		exit(EXIT_FAILURE);
	}
	//gpu_n = 1;
	int* num_gpu_blocks = (int *)malloc(gpu_n*sizeof(int));
	double** gpu_matrix_blocks =(double**)malloc(gpu_n * sizeof(a));
	double** gpu_column_blocks =(double**)malloc(gpu_n * sizeof(a));

	int pivot_row;
	int mkl_info;
	cublasHandle_t* handle = (cublasHandle_t *)malloc(gpu_n * sizeof(cublasHandle_t));
	cublasStatus_t stat;
	cudaStream_t* stream = (cudaStream_t *)malloc(gpu_n * sizeof(cudaStream_t));
	double *host_matrix_block = (double*)malloc(block_size*block_size*sizeof(double));

	//printf("Enabling P2P access\n");
	for(int i = 0; i < gpu_n; i++)
	{
	//	printf("Enabling %i device\n", i);
		checkCudaErrors(cudaSetDevice(i));
		for(int j = 0; (j < gpu_n) &&(j !=i); j++)
		{
			checkCudaErrors(cudaDeviceEnablePeerAccess(j, 0));
		}
	}
	
	int num_blocks;
	if(n%block_size)
	{
		num_blocks = n/block_size +1;
	} else
	{
		num_blocks = n/block_size;
	}

	int full_b = num_blocks/gpu_n;

	printf("Creating Streams\n");
	for(int i = 0; i < gpu_n; i++)
	{
		checkCudaErrors(cudaSetDevice(i));
		int curr_block_num = full_b;
		if(full_b*gpu_n + i < num_blocks)
		{
			curr_block_num++;
		}
		num_gpu_blocks[i] = curr_block_num;
		checkCudaErrors(cudaMalloc(&gpu_matrix_blocks[i], n* block_size * curr_block_num * sizeof(double)));
		checkCudaErrors(cudaMalloc(&gpu_column_blocks[i], n* block_size * sizeof(double)));
		
		checkCudaErrors(cudaStreamCreate(&stream[i]));
		
		stat = cublasCreate(&handle[i]);
		if(stat != CUBLAS_STATUS_SUCCESS)
		{
			printf("CUBLAS INITIALIZATION FAILED");
			return EXIT_FAILURE;
		}
		
		for(int j =0; j < curr_block_num; j++)
		{
			stat = cublasSetMatrixAsync(n, block_size, sizeof(double), &a[IDF2(0, j*gpu_n*block_size + i*block_size, n)], n, &gpu_matrix_blocks[i][j*n*block_size], n,
				stream[i]);
			if(stat != CUBLAS_STATUS_SUCCESS)
			{
				printf("%i %i\n", IDF2(0, j*gpu_n*block_size + i*block_size, n), j*n*block_size);
				printf("CAN NOT COPY BLOCK MATRIX, gpu %i and block %i, status %i ", i, j, stat);
				return EXIT_FAILURE;
			}
		}
	}
#ifdef DEBUG
	printf("Sync after copy\n");	
#endif
	for(int i = 0; i< gpu_n; i++)
	{
		checkCudaErrors(cudaSetDevice(i));
		cudaStreamSynchronize(stream[i]);
	}
	

	pivot_row = 0;
	mkl_info = 0;
	double alpha_1 = 1;
	double alpha_m1 = -1;
	int cur_gpu = -1;
	while(pivot_row + block_size < n)
	{
#ifdef DETAIL_DEBUG
		printf("\n\nStarting new iteration, pivot_row %i\n", pivot_row);
#endif
		cur_gpu = (pivot_row/block_size)%gpu_n;

		checkCudaErrors(cudaSetDevice(cur_gpu));
		checkCudaErrors(cudaDeviceSynchronize());
#ifdef DETAIL_DEBUG
		printf("!!!Current gpu %i blocks!!!!!!!!!\n", cur_gpu);
		cublasGetMatrix(n, num_gpu_blocks[i]*block_size, sizeof(double), gpu_matrix_blocks[cur_gpu], n, a, n);
		print_matrix(n, n, a, n);
		printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
#endif

		stat = cublasGetMatrix(block_size, block_size, sizeof(double), &(gpu_matrix_blocks[cur_gpu][IDF2(/*((pivot_row/block_size)/gpu_n)*block_size*/pivot_row, ((pivot_row/block_size)/gpu_n)*block_size, n)]), 
			n, host_matrix_block, block_size);
		if (stat != CUBLAS_STATUS_SUCCESS)
		{
			printf("CUBLAS DOWNLOAD MATRIX FAILED");
			return EXIT_FAILURE;
		}
#ifdef DETAIL_DEBUG
		printf("Start factorize block:\n");
		print_matrix(block_size, block_size, host_matrix_block, block_size);
#endif
		dpotrf("L", &block_size, host_matrix_block, &block_size, &mkl_info);
		if(mkl_info)
		{
			printf("BLOCK FACTORIZATION FAILED mkl_info - %d, pivot_row - %d \n", mkl_info, pivot_row);
			return EXIT_FAILURE;
		}
#ifdef DETAIL_DEBUG
		printf("End factorize matrix:\n");
		print_matrix(block_size, block_size, host_matrix_block, block_size);
#endif
		int block_column_gpu_number = (pivot_row/block_size)/gpu_n;
		int id_col_cur_gpu_block = ((pivot_row/block_size)/gpu_n)*block_size;
		cublasSetMatrix(block_size, block_size, sizeof(double), host_matrix_block, block_size, 
			&(gpu_matrix_blocks[cur_gpu][IDF2(/*((pivot_row/block_size)/gpu_n)*block_size*/pivot_row, id_col_cur_gpu_block, n)]), n);
#ifdef DETAIL_DEBUG
		printf("AFTER SET BLOCK\n");
		printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
		for(int q =0; q<n;q++)
			{
				for(int k =0; k <n;k++)
				{
					a[q*n+k] = 0;
				}
			}
		cublasGetMatrix(n, num_gpu_blocks[cur_gpu]*block_size, sizeof(double), gpu_matrix_blocks[cur_gpu], n, a, n);
		print_matrix(n, n, a, n);
		printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
#endif



		//find A21
#ifdef DETAIL_DEBUG
		printf("indexes (%i, %i) (%i, %i) \n", pivot_row, id_col_cur_gpu_block, 
			id_col_cur_gpu_block + cur_gpu*block_size + block_size, id_col_cur_gpu_block);
#endif
		stat = cublasDtrsm_v2(handle[cur_gpu], CUBLAS_SIDE_RIGHT, CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, n - pivot_row - block_size,
			block_size, &alpha_1, &(gpu_matrix_blocks[cur_gpu][IDF2(pivot_row, id_col_cur_gpu_block, n)]), n,
			&(gpu_matrix_blocks[cur_gpu][IDF2(pivot_row + block_size, 
			id_col_cur_gpu_block, n)]), n);

		if(stat != CUBLAS_STATUS_SUCCESS)
		{
			printf("CUBLA DTRSM ERROR %d %d", stat, (n - pivot_row - block_size));
			return EXIT_FAILURE;
		}
#ifdef DETAIL_DEBUG
		printf("AFTER FIND A21\n");
		printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
		cublasGetMatrix(n, num_gpu_blocks[cur_gpu]*block_size, sizeof(double), gpu_matrix_blocks[cur_gpu], n, a, n);
		print_matrix(n, n, a, n);
		printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
#endif

		
		//copy to gpu_colum_blocks
		for(int i = 0; (i < gpu_n) /*&& (i!=cur_gpu)*/; i++)
		{	
			//commented sync version
			//checkCudaErrors(cudaMemcpy(gpu_column_blocks[i], &gpu_matrix_blocks[cur_gpu][id_col_cur_gpu_block*n], n*block_size*sizeof(double), cudaMemcpyDefault ));
			//checkCudaErrors(cudaSetDevice(i));
			checkCudaErrors(cudaMemcpyAsync(gpu_column_blocks[i], &gpu_matrix_blocks[cur_gpu][id_col_cur_gpu_block*n], n*block_size*sizeof(double), cudaMemcpyDefault, stream[cur_gpu]));
		
		}

		checkCudaErrors(cudaStreamSynchronize(stream[cur_gpu]));
		/*for(int i =0; i< gpu_n; i++)
		{
			checkCudaErrors(cudaSetDevice(i));
			checkCudaErrors(cudaDeviceSynchronize());
			checkCudaErrors(cudaStreamSynchronize(stream[i]));
		}*/
		//update A22
		for(int i = 0; i < gpu_n ; i++)
		{	
			checkCudaErrors(cudaSetDevice(i));
			//checkCudaErrors(cudaDeviceSynchronize());
#ifdef DETAIL_DEBUG
			printf("!!!Column block for processing gpu - %i!!!!!\n", i);
			for(int q =0; q<n;q++)
			{
				for(int k =0; k <n;k++)
				{
					a[q*n+k] = 0;
				}
			}
			cublasGetMatrix(n, block_size, sizeof(double), gpu_column_blocks[i], n, a, n);
			print_matrix(n, n, a, n);
			printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
#endif
			
			int start_block = (pivot_row/block_size)/gpu_n;
			if(i <= cur_gpu)
			{
				start_block++;
			}
			
#ifdef DETAIL_DEBUG
			printf("start_block - %i, num_blocks - %i\n", start_block, num_gpu_blocks[i]);
#endif
			while(start_block < num_gpu_blocks[i])
			{
#ifdef DETAIL_DEBUG
				printf("!!!!!!!!!!!!!!!!! %i !!!!!!!!!!!!!\n", start_block);
				for(int q =0; q<n;q++)
				{
					for(int k =0; k <n;k++)
					{
						a[q*n+k] = 0;
					}
				}
				cublasGetMatrix(n - (start_block*gpu_n *block_size+ i*block_size), block_size, sizeof(double), 
					&(gpu_column_blocks[i][IDF2( start_block*gpu_n *block_size+ i*block_size, 0 ,n )]), n, a, n);
				print_matrix(n, n, a, n);
				printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");


				printf("!!!!!!!!!!!!!!!!! second %i !!!!!!!!!!!!!\n", start_block);
				for(int q =0; q<n;q++)
				{
					for(int k =0; k <n;k++)
					{
						a[q*n+k] = 0;
					}
				}
				cublasGetMatrix(block_size, block_size, sizeof(double), 
					&(gpu_column_blocks[i][IDF2( start_block*gpu_n *block_size+ i*block_size  ,0,n)]), n, a, n);
				print_matrix(n, n, a, n);
				printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
				printf("!!!!!!!!!!!!!!!!! third %i !!!!!!!!!!!!!\n", start_block);
				for(int q =0; q<n;q++)
				{
					for(int k =0; k <n;k++)
					{
						a[q*n+k] = 0;
					}
				}
				cublasGetMatrix(n - (start_block*gpu_n *block_size+ i*block_size), block_size, sizeof(double), 
					&(gpu_matrix_blocks[i][IDF2( start_block*gpu_n *block_size+ i*block_size , start_block*block_size, n )]), n, a, n);
				print_matrix(n, n, a, n);
				printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");

				printf("fact column start_block %i id %i\n", start_block, start_block*gpu_n *block_size+ i*block_size);
#endif
				int start_row_index = start_block*gpu_n *block_size+ i*block_size;

				stat = cublasDgemm_v2(handle[i], CUBLAS_OP_N, CUBLAS_OP_T, n - start_row_index, block_size, block_size,
					&alpha_m1, &(gpu_column_blocks[i][IDF2( start_row_index, 0 ,n )]), n, 
					&(gpu_column_blocks[i][IDF2( start_row_index ,0,n)]) , n,  &alpha_1,
					&(gpu_matrix_blocks[i][IDF2( start_row_index , start_block*block_size, n )]), n); 
				if (stat != CUBLAS_STATUS_SUCCESS)
				{
					printf("CUBLAS DGEMM FAILED");
					return EXIT_FAILURE;
				}
#ifdef DETAIL_DEBUG
				printf("AFTER ITER !!\n");
				printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
				cublasGetMatrix(n, num_gpu_blocks[i]*block_size, sizeof(double), gpu_matrix_blocks[i], n, a, n);
				print_matrix(n, n, a, n);
				printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
#endif
				start_block++;
			}
			//checkCudaErrors(cudaDeviceSynchronize());
#ifdef DETAIL_DEBUG
				printf("AFTER FIND A22 !!\n");
				printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
				cublasGetMatrix(n, num_gpu_blocks[i]*block_size, sizeof(double), gpu_matrix_blocks[i], n, a, n);
				print_matrix(n, n, a, n);
				printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
#endif
		}

		pivot_row+= block_size;

	}
	printf("Start factorize last block and copy back to host\n");
	for(int i = 0; i < gpu_n; i++)
	{
		checkCudaErrors(cudaSetDevice(i));
		int curr_block_num = full_b;
		if(full_b*gpu_n + i < num_blocks)
		{
			curr_block_num++;
		}
		for(int j =0; j < curr_block_num; j++)
		{
			stat = cublasGetMatrixAsync(n, block_size, sizeof(double), &gpu_matrix_blocks[i][j*n*block_size], n, &a[IDF2(0, j*gpu_n*block_size + i*block_size, n)], n,
				stream[i]);
			if(stat != CUBLAS_STATUS_SUCCESS)
			{
				printf("CAN NOT COPY BLOCK MATRIX\n");
				return EXIT_FAILURE;
			}

		}
	}
	printf("Sync after copy\n");
	
	for(int i = 0; i< gpu_n; i++)
	{
		checkCudaErrors(cudaSetDevice(i));
		cudaStreamSynchronize(stream[i]);
	}
#ifdef DETAIL_DEBUG
	printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
	print_matrix(n, n, a, n);
	printf("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n");
#endif
//factorize last block
	cur_gpu = (cur_gpu+1)%gpu_n;
	checkCudaErrors(cudaSetDevice(cur_gpu));
	stat = cublasGetMatrix(n - pivot_row, n - pivot_row, sizeof(double), &(gpu_matrix_blocks[cur_gpu][IDF2(/*((pivot_row/block_size)/gpu_n)*block_size*/pivot_row , 
		((pivot_row/block_size)/gpu_n)*block_size, n)]), n, host_matrix_block, block_size); 
		
	if (stat != CUBLAS_STATUS_SUCCESS)
	{
		printf("CUBLAS DOWNLOAD MATRIX FAILED, status - %i\n", stat);
		return EXIT_FAILURE;
	}
#ifdef DETAIL_DEBUG
	printf("Current pivot_row %i\n", pivot_row);
#endif
	int cur_block_size = n - pivot_row;
#ifdef DETAIL_DEBUG
	printf("LAST BLOCK TO SYNC %i  %i\n", block_size, cur_block_size);
	print_matrix(block_size, block_size, host_matrix_block, block_size);
	printf("=================\n");
#endif
	dpotrf("L", &cur_block_size, host_matrix_block, &block_size, &mkl_info);
	if(mkl_info)
	{
		printf("BLOCK FACTORIZATION FAILED %i, cur_block %i\n", mkl_info, cur_block_size);
		return EXIT_FAILURE;
	}
#ifdef DETAIL_DEBUG
	printf("LPLPLPL\n");
	print_matrix(n, n, a, n);
	printf("LPLPLPL\n");
#endif
	for(int i =0; i < cur_block_size; i++)
	{
		for(int j =0; j < cur_block_size; j++)
		{
			a[IDF2(pivot_row + i, pivot_row + j, n)] = host_matrix_block[IDF2(i,j, block_size)];
		}
	}
	return 0;
}
